package com.zdp.nio.channel.socketchannel.file;

import lombok.Builder;
import lombok.Data;
import lombok.SneakyThrows;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * @author zhengdp
 * @Description 文件传输 服务端 (用于接收文件）
 * @Date 2022/11/24 21:19
 * @Created by zhengdp
 */
@Data
public class SocketServer {
    @SneakyThrows
    public static void main(String[] args) {
        SocketServer server = new SocketServer(9905,"D:\\code\\");
        server.startServer();
    }

    public SocketServer(Integer port, String filePathPre) {
        this.port = port;
        this.filePathPre = filePathPre;
    }

    private Integer port;

    private Long length;

    private String filePathPre;

    private Charset charset = Charset.forName("UTF-8");

    private ByteBuffer buffer = ByteBuffer.allocate(1024);

    // 使用Map保存每个文件传输，当OP_READ可读时，根据通道找到对应的对象
    Map<SelectableChannel,Client> clientMap = new HashMap<>();

    public void startServer() throws IOException {
        // 1. 获取 选择器
        Selector selector = Selector.open();
        // 2. 创建 服务端监听通道
        ServerSocketChannel serverChannel = ServerSocketChannel.open();
        serverChannel.configureBlocking(false);
        // 2.1 绑定连接
        serverChannel.bind(new InetSocketAddress(port));
        // 2.2 注册到 selector
        serverChannel.register(selector,SelectionKey.OP_ACCEPT);
        // 3. selector 监听就绪的事件
        while(selector.select() > 0) {
            Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
            while(iterator.hasNext()) {
                SelectionKey next = iterator.next();
                if(next.isAcceptable()){
                    // 4. 处理 客户端连接请求事件
                    ServerSocketChannel server = (ServerSocketChannel) next.channel();
                    SocketChannel accept = server.accept();
                    accept.configureBlocking(false);
                    accept.register(selector,SelectionKey.OP_READ);
                    // 4.1 将对应的 客户端 Socket 与 Client 做映射
                    Client client = new Client();
                    client.setSocketAddress((InetSocketAddress) accept.getRemoteAddress());
                    clientMap.put(accept,client);
                    System.out.println(accept.getRemoteAddress() + " 连接成功....");
                }else if ( next.isReadable() ) {
                    // 读取文件
                    processData(next);
                }

                iterator.remove();
            }
        }
    }

    /**
     * 处理客户端传输过来的数据
     * 客户端传输的顺序 ：
     * 1. 文件名
     * 2. 文件长度
     * 3. 文件内容
     * @param key 通道事件
     * @throws IOException
     */
    private void processData(SelectionKey key) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        // 1. 获取到对应的 Client对象
        Client client = clientMap.get(channel);
        int num = 0;
        try {
            // 切换为写入模式
            buffer.clear();
            // 读取 通道中的数据到 Buffer对象
            while((num = channel.read(buffer)) > 0){
                while(buffer.position() != buffer.limit()) {
                    // 切换为读模式，从Buffer中读取数据，设置到内存对象中
                    buffer.flip();
                    if(length == null || length == 0){
                        length = buffer.getLong();
                    }
                    if(buffer.limit() == buffer.position()){
                        break;
                    }
                    byte[] bytes = new byte[1024];
                    if(null == client.fileName) {
                        // 设置文件名称 --> 连接的第一份数据，需要生成对应的文件通道，用于后续处
                        int bufferLength = buffer.limit() - buffer.position();
                        String fileName ;
                        if(bufferLength >= length){
                            buffer.get(bytes,buffer.position(),length.intValue());
                            fileName = new String(bytes,0,length.intValue(),charset);
                            length = 0L;
                        }else{
                            buffer.get(bytes,buffer.position(),buffer.limit());
                            fileName = new String(bytes,0,bufferLength,charset);
                            length -= bufferLength;
                        }
                        System.out.println("获取到的FileName : " + fileName);
                        File directory = new File(filePathPre);
                        if( ! directory.exists()){
                            directory.mkdir();
                        }
                        client.setFileName(fileName);
                        String fullFileName = directory.getAbsolutePath() + "\\" + fileName;
                        System.out.println("NIO 传输目标文件 ： " + fullFileName);
                        File file = new File(fullFileName);
                        FileChannel fileChannel = new FileOutputStream(file,true).getChannel();
                        client.setFileChannel(fileChannel);
                    }else if ( null == client.fileLength) {
                        System.out.println("处理长度");
                        int bufferLength = buffer.limit() - buffer.position();
                        if(bufferLength >= length){
                            buffer.get(bytes,buffer.position(),length.intValue());
                            length = 0L;
                        }else{
                            buffer.get(bytes,buffer.position(),buffer.limit());
                            length -= bufferLength;
                        }
                        // 设置文件 文件长度
                        long fileLength = ByteBuffer.wrap(bytes,0,8).getLong();
                        System.out.println("文件长度 ： " + fileLength);
                        client.setFileLength(fileLength);
                        client.startTime = System.currentTimeMillis();
                        System.out.println("NIO 传输开始 ： ");
                    }else {
                        // 写入文件内容
                        System.out.println("处理文件内容");
                        int bufferLength = buffer.limit() - buffer.position();
                        if(bufferLength >= length){
                            length = 0L;
                        }else{
                            length -= bufferLength;
                        }
                        while(buffer.hasRemaining()){
                            client.fileChannel.write(buffer);
                        }
                    }
                }
                buffer.clear();
            }
            // 事件处理完毕，废弃事件
            key.cancel();

        } catch (IOException e) {
            key.cancel();
            e.printStackTrace();
            return ;
        }

        // 读取数量-1，表示客户端传输结束标志到了
        if(num == -1){
            client.fileChannel.close();
            System.out.println("上传完毕");
            key.cancel();
            System.out.println("文件接收成功，File Name : " + client.fileName);
            System.out.println(" Size : " + client.fileLength);
            long endTime = System.currentTimeMillis();

            System.out.println("NIO IO 传输毫秒数 ： " + (endTime - client.startTime));
        }
    }

    @Data
    static class Client {
        // 文件名称
        String fileName;

        // 文件长度
        Long fileLength;

        // 开始传输时间
        Long startTime;

        // 客户端地址
        InetSocketAddress socketAddress;

        // 输出的文件通道
        FileChannel fileChannel;
    }
}
